from ultralytics import YOLO
import cv2
from ultralytics.yolo.utils.plotting import Annotator
import sys
sys.path.append("../")
from database.database import DataBase

class PlateDetect:
    def __init__(self, path=r"./ultralytics/weights/best.pt"):
        self.model = YOLO(path)
        self.color_list = [(239,71,112), (255,209,104), (6,216,161), (15,138,179), (6,58,75)]
        self.price_list = [1.0, 5.0, 5.0, 3.0, 5.0]
        # self.price_list = []
        # database = DataBase("./database/plave_data.db")
        # p_list = database.select_all("Price")
        # for i in database.select_all("Price"):
        #     print(i)
            # self.price_list.append(i[2])
        
        self.dish_name = ["rice", "dish1", "egg", "dish2", "fish"]

    def detect_frame(self, frame):
        results = self.model.predict(frame)
        dish_list = [{"rice": 0},{"dish1": 0},{"egg": 0},{"dish2": 0},{"fish": 0}]
        for r in results:
            annotator = Annotator(frame)
            boxes = r.boxes
            # print(self.model.names)
            total = 0
            for box in boxes:
                b = box.xyxy[0]
                c = box.cls
                conf = round(float(box.conf), 2)
                if conf < 0.8:
                    break
                annotator.box_label(b, f"{self.model.names[int(c)]}  {conf}", self.color_list[int(c)])
                # print(self.model.names[int(c)])
                # print(c)
                total += self.price_list[int(c)]
                dish_name = self.dish_name[int(c)]
                num = dish_list[int(c)][self.dish_name[int(c)]] + 1 
                # print(num)
                dish_list[int(c)][dish_name] = num

        frame = annotator.result()
        return frame, total, dish_list
    
def test_image():
    plate_engine = PlateDetect()
    print("模型加载完成")
    frame = cv2.imread("ultralytics/data/test_images/WIN_20211111_17_23_31_Pro.jpg")
    image, total, dish_list = plate_engine.detect_frame(frame)
    print(dish_list)
    print("图片推理成功")
    print(f"应付{total}元")
    cv2.imwrite("ultralytics/data/test3.jpg", image)

if __name__ == '__main__':
    test_image()
